import numpy as np


class ReservoirComputing:
    def __init__(self, input_dim, reservoir_dim, output_dim, spectral_radius=0.95, sparsity=0.1, alpha=1.0):
        self.input_dim = input_dim
        self.reservoir_dim = reservoir_dim
        self.output_dim = output_dim
        self.spectral_radius = spectral_radius
        self.sparsity = sparsity
        self.alpha = alpha

        # 初始化权重矩阵
        self.Win = np.random.rand(self.reservoir_dim, self.input_dim) - 0.5
        self.W = np.random.rand(self.reservoir_dim, self.reservoir_dim) - 0.5

        # 稀疏化
        mask = np.random.rand(self.reservoir_dim, self.reservoir_dim) < self.sparsity
        self.W[mask] = 0

        # 调整谱半径
        radius = np.max(np.abs(np.linalg.eigvals(self.W)))
        self.W *= self.spectral_radius / radius

        self.Wout = None
        self.states = np.zeros((self.reservoir_dim, 1))

    def _activation(self, x):
        return np.tanh(x)

    def train(self, input_data, target_data):
        num_samples = input_data.shape[0]
        states_collection = np.zeros((num_samples, self.reservoir_dim))

        # 驱动储备池
        for t in range(num_samples):
            u = input_data[t]
            self.states = (1 - self.alpha) * self.states + self.alpha * self._activation(
                self.Win @ u + self.W @ self.states)
            states_collection[t] = self.states[:, 0]

        # 增加偏置
        states_collection = np.hstack((states_collection, np.ones((num_samples, 1))))

        # 训练读取层（线性回归）
        self.Wout = np.linalg.pinv(states_collection) @ target_data

    def predict(self, input_data):
        num_samples = input_data.shape[0]
        predictions = np.zeros((num_samples, self.output_dim))

        # 驱动储备池并进行预测
        for t in range(num_samples):
            u = input_data[t]
            self.states = (1 - self.alpha) * self.states + self.alpha * self._activation(
                self.Win @ u + self.W @ self.states)
            extended_state = np.hstack((self.states[:, 0], [1]))  # 增加偏置
            predictions[t] = self.Wout @ extended_state

        return predictions


# 示例使用
if __name__ == "__main__":
    # 生成一些示例数据
    t = np.linspace(0, 10, 100)
    input_data = np.sin(t).reshape(-1, 1)
    target_data = np.cos(t).reshape(-1, 1)

    # 初始化储备池计算
    rc = ReservoirComputing(input_dim=1, reservoir_dim=100, output_dim=1)

    # 训练模型
    rc.train(input_data, target_data)

    # 预测
    predictions = rc.predict(input_data)

    # 打印结果
    print(predictions)
